#!/usr/bin/python

import sys, math, random
import numpy as np
from profilehooks import profile
import bisect as bi
import ss_hash as ss
import util


# integrate with space saving. Do lightweight implementation. Min.Heap
# change iteration to work with DCSC, DCSR
# implement MPI / mmap
# exploit symmetry: upper/lower triangular
# modulo. http://www.geeksforgeeks.org/archives/9057 . using 2^n processors
# random sampling. is it prob v'[i]/sum(v) i in [a:b) or

p = 10 # number of processors

def sample(prob):
    return  bi.bisect_left(prob, random.random()) # returns insertion index

def generate_matrix(n, r) :
    global A,B, C_mat, C, p_dist, C, h1, h2
    C, h1, h2 = {}, {},{}
    p_dist = np.zeros((p,), dtype=np.int)
    A = np.random.randint(r, size=(n, n))
    B = np.random.randint(r, size=(n, n))
    C_mat = [[0 for i in range(n)] for x in range(n)]

#@profile
def val(h, x, hi):
    if not(x in h):
	random.seed(x + hi) # http://docs.python.org/library/random.htm. Hashvalue of x is determined by the random.seed
	h[x] = int(random.random() * p)
    return h[x]

#@profile
def generate_prob(h1_counts, h1val) :
    global prob_ranges
    prob_ranges = []

    for pi in range(len(h1_counts)-1) :
        # get range and compute accumulated probability
        ab_range = h1val[h1_counts[pi]:h1_counts[pi + 1]]
        prob = list(map(lambda x: x[2], ab_range)) # get list of values
        ab_sum = sum(prob)
        prob = list( map(lambda x : x / ab_sum, np.cumsum(prob))) # normalize accumulated weights
        prob_ranges.append(prob)

#@profile
def do_outer_product(h1val, h2val, h1_counts) :
    for pi in range(p) :
    	for b_item in h2val :
            #print "size", int (len(A) * b_item[2] / 10 + 0.5)
            SS = ss.StreamSummary(int (len(A)) )

    	    a_idx = (p + pi - b_item[0]) % p # calculate the value of "a" for which (h1(a_i) + h2(b_j)) % p = pi
    	    ab_range = h1val[h1_counts[a_idx]:h1_counts[a_idx + 1]] # get range of items from A
            if len(ab_range) == 0 : continue # processor pi not responsible for this multiplication

    	    p_dist[pi] += len(ab_range) # debugging

            prob = prob_ranges[a_idx]
            for x in range(int(math.ceil(len(A) / 4))) : # how to determine the appropriate number of samples
                r = sample(prob)
                SS.insert(r)

            for key in SS.elements.keys() : # iterate saved elements in the space saving sketch
                a_item = ab_range[key]
                key = (a_item[1], b_item[1])
                if key not in C :
                    C[key] = 0
                C[key] += b_item[2] * a_item[2]

#@profile
def CRoP() :
    for a_k, col in enumerate(A.T) :
    	h1_counts = [0 for x in range(p)] # accumulated counts for hashed col indices of A
    	h1val, h2val, = [], []
        # h1val, h2val: list of tuples: (h1(a_i), a_i, col[a_i]),

    	accum_sum = float(sum(col)) # sum of values in col a_k
    	if accum_sum == 0 :
            continue

    	for a_i, item in enumerate(col) :
    	    a_val = float(item)
    	    if a_val == 0 : continue

    	    hash1 = val(h1, a_i, 0)
    	    h1_counts[hash1] += 1
    	    h1val.append((hash1, a_i, a_val)) # insert tuple into h1val

    	for b_j, value in enumerate(B[a_k]) :
    	    if value  == 0 : continue
    	    h2val.append((val(h2, b_j, 1), b_j, float(value))) # insert tuple into h2val

    	h1_counts = list(np.cumsum(h1_counts))
        #print "h1_counts", h1_counts
        h1_counts.insert(0,0)
        h1_counts.sort()
        generate_prob(h1_counts, h1val) # precompute probability distribution of each range in h1val
    	do_outer_product(h1val, h2val, h1_counts) # do outer products




def matrix_mult(A, B) :
    return A * B

generate_matrix(100, 4)
#print (A)
#print (B)

def callback(*args, **kwargs):
    print (args, kwargs)

with util.Timer() as t :
    CRoP()
print('CRoP took %.03f sec.' % t.interval)

# build result matrix from C
for key, val in C.iteritems():
    try :
        C_mat[key[0]][key[1]] = val
    except Exception as e :
        print e
        print key, val

print "p_dist: ", p_dist
print "avg: ", np.average(p_dist)
print "mean: ", np.mean(p_dist)
print "std: ", np.std(p_dist)

A = np.mat(A)
B = np.mat(B)
COUT = matrix_mult(A,B)
C_mat = np.mat(C_mat)

print (np.array_equal(COUT, C_mat))